"""
程序入口模块

该模块作为程序的主入口点，负责：
1. 解析命令行参数
2. 加载配置文件
3. 初始化计算环境
4. 执行计算流程
5. 输出结果
"""

import argparse
import json
import numpy as np
from pathlib import Path
from typing import Dict, Optional

from src.geometry.subchannel import SubchannelGeometry
from src.core.staggered_grid import StaggeredGrid
from src.core.solver import ConservationSolver
from src.core.numerical import NumericalMethods
from src.physics.two_phase import TwoPhaseFlow
from src.utils.logger import setup_logger
from src.utils.visualization import ResultVisualizer

class SubchannelAnalysis:
    """子通道分析主程序类"""
    
    def __init__(self, config_file: str):
        """
        初始化子通道分析程序
        
        参数:
            config_file (str): 配置文件路径
        """
        # 读取配置文件
        with open(config_file, 'r') as f:
            self.config = json.load(f)
            
        # 初始化几何模型
        self.geometry = SubchannelGeometry(self.config)
        
        # 基于几何信息初始化网格
        self.grid = StaggeredGrid(
            
            self.config
        )
        
        # 初始化数值方法
        self.numerical = NumericalMethods(self.config)
        
        # 初始化两相流模型
        self.two_phase = TwoPhaseFlow(self.grid)
        
        # 初始化求解器
        self.solver = ConservationSolver(
            self.grid,
            self.numerical,
            self.two_phase,
            self.config
        )
        
        # 初始化结果存储
        self.results = {
            'time': [],
            'pressure': [],
            'temperature': [],
            'void_fraction': [],
            'velocity': []
        }
        
    def run(self):
        """运行分析"""
        # 获取时间步长和总时间
        dt = self.config['numerical']['time_step']
        t_max = self.config['numerical']['max_time']
        save_interval = self.config['numerical']['save_interval']
        
        # 时间推进循环
        t = 0.0
        next_save = save_interval
        while t < t_max:
            # 求解当前时间步
            converged = self.solver.solve_time_step(dt)
            
            if not converged:
                print(f"警告: 时间步 {t:.3f} s 未收敛")
                
            # 保存结果
            if t >= next_save:
                self._save_results(t)
                next_save += save_interval
                
            # 更新时间
            t += dt
            
        # 保存最终结果
        self._save_results(t)
        
    def _save_results(self, t: float):
        """
        保存当前时间步的结果
        
        参数:
            t (float): 当前时间
        """
        self.results['time'].append(t)
        self.results['pressure'].append(
            self.grid.get_scalar_field('pressure').copy()
        )
        self.results['temperature'].append(
            self.grid.get_scalar_field('temperature').copy()
        )
        self.results['void_fraction'].append(
            self.grid.get_scalar_field('void_fraction').copy()
        )
        self.results['velocity'].append({
            'axial': self.grid.get_axial_velocity_field().copy(),
            'lateral': self.grid.get_lateral_velocity_field().copy()
        })
        
    def save_to_file(self, output_file: str):
        """
        将结果保存到文件
        
        参数:
            output_file (str): 输出文件路径
        """
        # 转换结果为numpy数组
        results_np = {
            'time': np.array(self.results['time']),
            'pressure': np.array(self.results['pressure']),
            'temperature': np.array(self.results['temperature']),
            'void_fraction': np.array(self.results['void_fraction']),
            'velocity_axial': np.array([v['axial'] for v in self.results['velocity']]),
            'velocity_lateral': np.array([v['lateral'] for v in self.results['velocity']])
        }
        
        # 保存到文件
        np.savez(output_file, **results_np)

def main():
    """主函数"""
    # 设置日志
    logger = setup_logger(log_file="subchannel.log")
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='核反应堆子通道分析程序')
    parser.add_argument('config', help='配置文件路径')
    parser.add_argument('-o', '--output', help='输出文件路径', default='results.npz')
    parser.add_argument('-p', '--plot', help='结果图片保存目录', default='plots')
    args = parser.parse_args()
    
    try:
        # 创建并运行分析程序
        logger.info("开始运行子通道分析程序")
        analysis = SubchannelAnalysis(args.config)
        analysis.run()
        
        # 保存结果
        logger.info(f"保存结果到文件: {args.output}")
        analysis.save_to_file(args.output)
        
        # 可视化结果
        logger.info(f"生成结果可视化图片: {args.plot}")
        visualizer = ResultVisualizer(analysis.results)
        visualizer.plot_axial_profiles(args.plot)
        
        logger.info("程序运行完成")
        
    except Exception as e:
        logger.error(f"程序运行出错: {str(e)}")
        raise

if __name__ == '__main__':
    main() 